﻿Imports SlimDX
Imports SlimDX.Direct3D9


Public Class ShaderGeneric
    Implements IShader

    Private MapBlock As Array
    Private MapCount As Integer

    Private StageBlock As Array
    Private StageCount As Integer

    Private Maps() As Texture

    Private sotrEffect As Effect

    Dim Flags, FramebufferFunction As Short
    Sub New(ByVal sotrTag As Magic.Core.Tag)
        AddHandler sotrTag.TagHasChanged, AddressOf Changed
        Flags = sotrTag(13)
        FramebufferFunction = sotrTag(15)

        MapBlock = sotrTag("maps")
        MapCount = MapBlock.GetLength(0)

        StageBlock = sotrTag("stages")
        StageCount = StageBlock.GetLength(0)

        Dim EffectString As String = GenerateEffect()

        sotrEffect = Effect.FromString(xDevice, EffectString, Nothing, Nothing, ShaderFlags.None, Nothing)
        ReDim Maps(MapCount - 1)
        For x As Integer = 0 To MapCount - 1
            Dim MapIndex As Integer = MapBlock.GetValue(x, 8)
            If MapIndex <> -1 Then
                Maps(x) = New Bitmap(magicMap.Library(MapIndex)).Texture
            Else
                Maps(x) = Globals.default2d(IShader.DefaultTextures.Additive)
            End If
            sotrEffect.SetTexture("Map" & x.ToString, Maps(x))
            sotrEffect.SetValue("map" & x.ToString & "Scale", New Vector4(MapBlock.GetValue(x, 2), MapBlock.GetValue(x, 3), 0, 0))
        Next
        Dim constant0Low(6) As Color4
        Dim constant0High(6) As Color4
        Dim constant1(6) As Color4
        For x As Integer = 0 To StageCount - 1
            constant0Low(x) = New Color4(StageBlock.GetValue(x, 6)(1), StageBlock.GetValue(x, 6)(2), StageBlock.GetValue(x, 6)(3))
            constant0High(x) = New Color4(StageBlock.GetValue(x, 7)(1), StageBlock.GetValue(x, 7)(2), StageBlock.GetValue(x, 7)(3))
            constant1(x) = New Color4(StageBlock.GetValue(x, 8)(1), StageBlock.GetValue(x, 8)(2), StageBlock.GetValue(x, 8)(3))
        Next
        sotrEffect.SetValue("constantColor0Low", constant0Low)
        sotrEffect.SetValue("constantColor0High", constant0High)
        sotrEffect.SetValue("constantColor1", constant1)

        sotrEffect.Technique = "GenericTransparent"
    End Sub
    Sub Changed(ByVal sotrTag As Magic.Core.Tag)
        Flags = sotrTag(13)
        FramebufferFunction = sotrTag(15)

        MapBlock = sotrTag("maps")
        MapCount = MapBlock.GetLength(0)

        StageBlock = sotrTag("stages")
        StageCount = StageBlock.GetLength(0)

        Dim EffectString As String = GenerateEffect()

        sotrEffect = Effect.FromString(xDevice, EffectString, Nothing, Nothing, ShaderFlags.None, Nothing)
        ReDim Maps(MapCount - 1)
        For x As Integer = 0 To MapCount - 1
            Dim MapIndex As Integer = MapBlock.GetValue(x, 8)
            If MapIndex <> -1 Then
                Maps(x) = New Bitmap(magicMap.Library(MapIndex)).Texture
            Else
                Maps(x) = Globals.default2d(IShader.DefaultTextures.Additive)
            End If
            sotrEffect.SetTexture("Map" & x.ToString, Maps(x))
            sotrEffect.SetValue("map" & x.ToString & "Scale", New Vector4(MapBlock.GetValue(x, 2), MapBlock.GetValue(x, 3), 0, 0))
        Next
        Dim constant0Low(6) As Color4
        Dim constant0High(6) As Color4
        Dim constant1(6) As Color4
        For x As Integer = 0 To StageCount - 1
            constant0Low(x) = New Color4(StageBlock.GetValue(x, 6)(1), StageBlock.GetValue(x, 6)(2), StageBlock.GetValue(x, 6)(3))
            constant0High(x) = New Color4(StageBlock.GetValue(x, 7)(1), StageBlock.GetValue(x, 7)(2), StageBlock.GetValue(x, 7)(3))
            constant1(x) = New Color4(StageBlock.GetValue(x, 8)(1), StageBlock.GetValue(x, 8)(2), StageBlock.GetValue(x, 8)(3))
        Next
        sotrEffect.SetValue("constantColor0Low", constant0Low)
        sotrEffect.SetValue("constantColor0High", constant0High)
        sotrEffect.SetValue("constantColor1", constant1)

        sotrEffect.Technique = "GenericTransparent"
    End Sub
    Public Overloads Sub BeginPass(ByVal pass As Integer) Implements IShader.BeginPass
        If (pass = 0) Then
            xDevice.SetRenderState(RenderState.AlphaBlendEnable, True)
            xDevice.SetRenderState(RenderState.AlphaTestEnable, (Flags And 1 <> 0))

            xDevice.SetRenderState(RenderState.AlphaRef, &H80)
            xDevice.SetRenderState(RenderState.AlphaFunc, Compare.GreaterEqual)

            xDevice.SetRenderState(RenderState.CullMode, IIf(Flags And 4 <> 0, Cull.None, Cull.Counterclockwise))
            SetFramebufferFunction(FramebufferFunction)
        End If
        sotrEffect.BeginPass(pass)
    End Sub
    Public Overloads Sub EndPass() Implements IShader.EndPass
        sotrEffect.EndPass()
    End Sub
    Public Overloads Function BeginShader(ByVal tLightmap As Texture) As Integer Implements IShader.BeginShader
        'Globals
        sotrEffect.SetValue("cameraPosition", New Vector4(Camera.Position.X, Camera.Position.Y, Camera.Position.Z, 1.0F))
        sotrEffect.SetValue("world", xDevice.GetTransform(TransformState.World))
        sotrEffect.SetValue("worldViewProjection", xDevice.GetTransform(TransformState.World) * xDevice.GetTransform(TransformState.View) * xDevice.GetTransform(TransformState.Projection))
        sotrEffect.SetValue("fogStart", xDevice.GetRenderState(RenderState.FogStart)) '0.0F) ' MdxRender.Device.RenderState.FogStart);
        sotrEffect.SetValue("fogEnd", xDevice.GetRenderState(RenderState.FogEnd)) ' MdxRender.Device.RenderState.FogEnd);
        sotrEffect.SetValue("fogDensity", xDevice.GetRenderState(RenderState.FogDensity)) ' MdxRender.Device.RenderState.FogDensity);

        Dim colorValues(6) As Single
        For x As Integer = 0 To StageCount - 1
            colorValues(x) = Animate(StageBlock.GetValue(x, 4), StageBlock.GetValue(x, 5), 1)
        Next
        sotrEffect.SetValue("colorValues", colorValues)

        For x As Integer = 0 To MapCount - 1
            sotrEffect.SetValue("map" & x.ToString & "Offset", New Vector4(MapBlock.GetValue(x, 4) + Animate(MapBlock.GetValue(x, 11), MapBlock.GetValue(x, 12), 1) * MapBlock.GetValue(x, 14), MapBlock.GetValue(x, 5) + Animate(MapBlock.GetValue(x, 16), MapBlock.GetValue(x, 17), 1) * MapBlock.GetValue(x, 19), 1, 1))
        Next

        Return sotrEffect.Begin(FX.DoNotSaveState)
    End Function
    Public Overloads Sub EndShader() Implements IShader.EndShader
        sotrEffect.End()
    End Sub
    Private Sub SetFramebufferFunction(ByVal [function] As Short)
        Select Case [function]
            Case 0 ': // alpha blend
                xDevice.SetRenderState(RenderState.BlendOperation, BlendOperation.Add)
                xDevice.SetRenderState(RenderState.SourceBlend, Blend.SourceAlpha)
                xDevice.SetRenderState(RenderState.DestinationBlend, Blend.InverseSourceAlpha)
            Case 1, 2 ': // multiply
                xDevice.SetRenderState(RenderState.BlendOperation, BlendOperation.Add)
                xDevice.SetRenderState(RenderState.SourceBlend, Blend.SourceColor)
                xDevice.SetRenderState(RenderState.DestinationBlend, Blend.DestinationColor)
            Case 3 ': // add
                xDevice.SetRenderState(RenderState.BlendOperation, BlendOperation.Add)
                xDevice.SetRenderState(RenderState.SourceBlend, Blend.One)
                xDevice.SetRenderState(RenderState.DestinationBlend, Blend.One)
            Case 4 ': // subtract
                xDevice.SetRenderState(RenderState.BlendOperation, BlendOperation.Subtract)
                xDevice.SetRenderState(RenderState.SourceBlend, Blend.One)
                xDevice.SetRenderState(RenderState.DestinationBlend, Blend.One)
            Case 5 ': // component min
                xDevice.SetRenderState(RenderState.BlendOperation, BlendOperation.Minimum)
                xDevice.SetRenderState(RenderState.SourceBlend, Blend.One)
                xDevice.SetRenderState(RenderState.DestinationBlend, Blend.One)
            Case 6 ': // component max
                xDevice.SetRenderState(RenderState.BlendOperation, BlendOperation.Maximum)
                xDevice.SetRenderState(RenderState.SourceBlend, Blend.One)
                xDevice.SetRenderState(RenderState.DestinationBlend, Blend.One)
            Case 7 ': // alpha-multiply add
                xDevice.SetRenderState(RenderState.BlendOperation, BlendOperation.Add)
                xDevice.SetRenderState(RenderState.SourceBlend, Blend.SourceAlpha)
                xDevice.SetRenderState(RenderState.DestinationBlend, Blend.DestinationAlpha)
        End Select
    End Sub

    Private Function GenerateEffect() As String
        Dim output As New System.Text.StringBuilder
        output.AppendFormat("#define MAX_STAGE_COUNT {0}" & Environment.NewLine, 7)
        For x As Integer = 0 To MapCount - 1
            If CInt(MapBlock.GetValue(x, 0)) And 1 = 1 Then
                output.AppendFormat("#define UNFILTERED_{0} 1" & Environment.NewLine, x)
            End If
        Next
        output.Append(System.Text.ASCIIEncoding.ASCII.GetString(My.Resources.generic_a))

        If StageCount = 0 Then
            output.Append(CodeLine("scratch_0 = map_0"))
        Else
            For i As Integer = 0 To StageCount - 1
                output.AppendFormat("// stage {0}" & Environment.NewLine, i)

                output.AppendFormat(CodeLine("constant_0 = lerp(constantColor0Low[{0}], constantColor0High[{0}], colorValues[{0}])"), i)
                output.AppendFormat(CodeLine("constant_1 = constantColor1[{0}]"), i)

                output.AppendFormat(CodeLine("alphaInputA = {0}"), InputMappingString(AlphaInputSourceString(StageBlock.GetValue(i, 26)), StageBlock.GetValue(i, 27)))
                output.AppendFormat(CodeLine("alphaInputB = {0}"), InputMappingString(AlphaInputSourceString(StageBlock.GetValue(i, 28)), StageBlock.GetValue(i, 29)))
                output.AppendFormat(CodeLine("alphaInputC = {0}"), InputMappingString(AlphaInputSourceString(StageBlock.GetValue(i, 30)), StageBlock.GetValue(i, 31)))
                output.AppendFormat(CodeLine("alphaInputD = {0}"), InputMappingString(AlphaInputSourceString(StageBlock.GetValue(i, 32)), StageBlock.GetValue(i, 33)))
                output.Append(CodeLine("alphaOutputAB = alphaInputA * alphaInputB"))
                output.Append(CodeLine("alphaOutputCD = alphaInputC * alphaInputD"))
                If (StageBlock.GetValue(i, 0) And 2 = 2) Then '// alpha mux
                    output.Append(CodeLine("alphaOutputABCD = (scratch_0.a > 0.5f) ? alphaOutputAB : alphaOutputCD"))
                Else
                    output.Append(CodeLine("alphaOutputABCD = alphaOutputAB + alphaOutputCD"))
                End If
                If (StageBlock.GetValue(i, 35) <> 0) Then
                    output.AppendFormat(CodeLine("{0}.a = alphaOutputAB"), OutputDestinationString(StageBlock.GetValue(i, 35)))
                End If
                If (StageBlock.GetValue(i, 36) <> 0) Then
                    output.AppendFormat(CodeLine("{0}.a = alphaOutputCD"), OutputDestinationString(StageBlock.GetValue(i, 36)))
                End If
                If (StageBlock.GetValue(i, 37) <> 0) Then
                    output.AppendFormat(CodeLine("{0}.a = alphaOutputABCD"), OutputDestinationString(StageBlock.GetValue(i, 37)))
                End If
                If (StageBlock.GetValue(i, 0) And 4 = 4) Then ' // a-out controls color0 animation
                    output.AppendFormat(CodeLine("constant_0 = lerp(constantColor0Low[{0}], constantColor0High[{0}], alphaOutputABCD)"), i)
                End If
                output.AppendFormat(CodeLine("colorInputA = {0}"), InputMappingString(ColorInputSourceString(StageBlock.GetValue(i, 10)), StageBlock.GetValue(i, 11)))
                output.AppendFormat(CodeLine("colorInputB = {0}"), InputMappingString(ColorInputSourceString(StageBlock.GetValue(i, 12)), StageBlock.GetValue(i, 13)))
                output.AppendFormat(CodeLine("colorInputC = {0}"), InputMappingString(ColorInputSourceString(StageBlock.GetValue(i, 14)), StageBlock.GetValue(i, 15)))
                output.AppendFormat(CodeLine("colorInputD = {0}"), InputMappingString(ColorInputSourceString(StageBlock.GetValue(i, 16)), StageBlock.GetValue(i, 17)))
                output.Append(CodeLine("colorOutputAB = " + OutputFunctionString("colorInputA", "colorInputB", StageBlock.GetValue(i, 20))))
                output.Append(CodeLine("colorOutputCD = " + OutputFunctionString("colorInputC", "colorInputD", StageBlock.GetValue(i, 22))))
                If (StageBlock.GetValue(i, 0) And 1 = 1) Then '// color mux
                    output.Append(CodeLine("colorOutputABCD = (scratch_0.a > 0.5f) ? colorOutputAB : colorOutputCD"))
                Else
                    output.Append(CodeLine("colorOutputABCD = colorOutputAB + colorOutputCD"))
                End If
                If (StageBlock.GetValue(i, 19) <> 0) Then
                    output.AppendFormat(CodeLine("{0}.rgb = colorOutputAB.rgb"), OutputDestinationString(StageBlock.GetValue(i, 19)))
                End If
                If (StageBlock.GetValue(i, 21) <> 0) Then
                    output.AppendFormat(CodeLine("{0}.rgb = colorOutputCD.rgb"), OutputDestinationString(StageBlock.GetValue(i, 21)))
                End If
                If (StageBlock.GetValue(i, 23) <> 0) Then
                    output.AppendFormat(CodeLine("{0}.rgb = colorOutputABCD.rgb"), OutputDestinationString(StageBlock.GetValue(i, 23)))
                End If
                output.Append(Environment.NewLine)
            Next
        End If

        output.Append(System.Text.ASCIIEncoding.ASCII.GetString(My.Resources.generic_b))
        Return output.ToString()
    End Function
    Private Function OutputDestinationString(ByVal destination As Short) As String

        If (destination = 1) Then '// scratch color 0
            Return "scratch_0"
        ElseIf (destination = 2) Then '// scratch color 1
            Return "scratch_1"
        ElseIf (destination = 3) Then '// vertex color 0
            Return "vertex_0"
        ElseIf (destination = 4) Then '// vertex color 1
            Return "vertex_1"
        ElseIf (destination = 5) Then '// map color 0
            Return "map_0"
        ElseIf (destination = 6) Then '// map color 1
            Return "map_1"
        ElseIf (destination = 7) Then '// map color 2
            Return "map_2"
        ElseIf (destination = 8) Then '// map color 3
            Return "map_3"
        End If
        Return Nothing
        'throw new HaloException("Could not map output destination field.");
    End Function
    Private Function OutputFunctionString(ByVal a As String, ByVal b As String, ByVal [function] As Short) As String
        If ([function] = 0) Then '// multiply
            Return a + " * " + b
        ElseIf ([function] = 1) Then '// dot product
            Return "dot(" + a + ", " + b + ")"
        End If
        Return Nothing
        'throw new HaloException("Failed to map an output function in a generic shader.");
    End Function
    Private Function InputMappingString(ByVal input As String, ByVal mapping As Short) As String
        If (mapping = 0) Then '// clamp(x)
            Return "saturate(" + input + ")"
        ElseIf (mapping = 1) Then '// 1 - clamp(x)
            Return "1.0f - saturate(" + input + ")"
        ElseIf (mapping = 2) Then '// 2 * clamp(x) - 1
            Return "2.0f * saturate(" + input + ") - 1.0f"
        ElseIf (mapping = 3) Then '// 1 - 2 * clamp(x)
            Return "1.0f - 2.0f * saturate(" + input + ")"
        ElseIf (mapping = 4) Then '// clamp(x) - 1/2
            Return "saturate(" + input + ") - 0.5f"
        ElseIf (mapping = 5) Then '// 1/2 - clamp(x)
            Return "0.5f - saturate(" + input + ")"
        ElseIf (mapping = 6) Then '// x
            Return input
        ElseIf (mapping = 7) Then '// -x
            Return "-" + input
        End If
        Return Nothing
        ' throw new HaloException("Failed to map an input function in a generic shader.");
    End Function
    Private Function AlphaInputSourceString(ByVal source As Short) As String
        If (source = 0) Then ' zero
            Return "0.0f"
        ElseIf (source = 1) Then ' one
            Return "1.0f"
        ElseIf (source = 2) Then ' one half
            Return "0.5f"
        ElseIf (source = 3) Then ' negative one
            Return "-1.0f"
        ElseIf (source = 4) Then ' negative one half
            Return "-0.5f"
        ElseIf (source = 5) Then ' map alpha 0
            Return "map_0.a"
        ElseIf (source = 6) Then ' map alpha 1
            Return "map_1.a"
        ElseIf (source = 7) Then ' map alpha 2
            Return "map_2.a"
        ElseIf (source = 8) Then ' map alpha 3
            Return "map_3.a"
        ElseIf (source = 9) Then ' vertex alpha 0
            Return "vertex_0.a"
        ElseIf (source = 10) Then ' vertex alpha 1
            Return "vertex_1.a"
        ElseIf (source = 11) Then ' scratch alpha 0
            Return "scratch_0.a"
        ElseIf (source = 12) Then ' scratch alpha 1
            Return "scratch_1.a"
        ElseIf (source = 13) Then ' constant alpha 0
            Return "constant_0.a"
        ElseIf (source = 14) Then ' constant alpha 1
            Return "constant_1.a"
        ElseIf (source = 15) Then ' map blue 0
            Return "map_0.b"
        ElseIf (source = 16) Then ' map blue 1
            Return "map_1.b"
        ElseIf (source = 17) Then ' map blue 2
            Return "map_2.b"
        ElseIf (source = 18) Then ' map blue 3
            Return "map_3.b"
        ElseIf (source = 19) Then ' vertex blue 0
            Return "vertex_0.b"
        ElseIf (source = 20) Then ' vertex blue 1
            Return "vertex_1.b"
        ElseIf (source = 21) Then ' scratch blue 0
            Return "scratch_0.b"
        ElseIf (source = 22) Then ' scratch blue 1
            Return "scratch_1.b"
        ElseIf (source = 23) Then ' constant blue 0
            Return "constant_0.b"
        ElseIf (source = 24) Then ' constant blue 1
            Return "constant_1.b"
        End If
        Return Nothing
        'throw new HaloException("Could not generate alpha input argument.");
    End Function
    Private Function ColorInputSourceString(ByVal source As Short) As String
        If (source = 0) Then '// zero
            Return "0.0f"
        ElseIf (source = 1) Then '// one
            Return "1.0f"
        ElseIf (source = 2) Then '// one half
            Return "0.5f"
        ElseIf (source = 3) Then '// negative one
            Return "-1.0f"
        ElseIf (source = 4) Then '// negative one half
            Return "-0.5f"
        ElseIf (source = 5) Then '// map 0
            Return "map_0"
        ElseIf (source = 6) Then '// map 1
            Return "map_1"
        ElseIf (source = 7) Then '// map 2
            Return "map_2"
        ElseIf (source = 8) Then '// map 3
            Return "map_3"
        ElseIf (source = 9) Then '// vertex 0
            Return "vertex_0"
        ElseIf (source = 10) Then '// vertex 1
            Return "vertex_1"
        ElseIf (source = 11) Then '// scratch 0
            Return "scratch_0"
        ElseIf (source = 12) Then '// scratch 1
            Return "scratch_1"
        ElseIf (source = 13) Then '// constant 0
            Return "constant_0"
        ElseIf (source = 14) Then '// constant 1
            Return "constant_1"
        ElseIf (source = 15) Then '// map alpha 0
            Return "map_0.a"
        ElseIf (source = 16) Then '// map alpha 1
            Return "map_1.a"
        ElseIf (source = 17) Then '// map alpha 2
            Return "map_2.a"
        ElseIf (source = 18) Then '// map alpha 3
            Return "map_3.a"
        ElseIf (source = 19) Then '// vertex alpha 0
            Return "vertex_0.a"
        ElseIf (source = 20) Then '// vertex alpha 1
            Return "vertex_1.a"
        ElseIf (source = 21) Then '// scratch alpha 0
            Return "scratch_0.a"
        ElseIf (source = 22) Then '// scratch alpha 1
            Return "scratch_1.a"
        ElseIf (source = 23) Then '// constant alpha 0
            Return "constant_0.a"
        ElseIf (source = 24) Then '// constant alpha 1
            Return "constant_1.a"
        End If
        Return Nothing
        'throw new HaloException("Could not generate color input argument.");
    End Function
    Private Shared Function CodeLine(ByVal input As String) As String
        Return vbTab & input & ";" & Environment.NewLine
    End Function
End Class
